from ultralytics import YOLO
import torch
import os
# 设置环境变量，避免某些库的冲突
os.environ['KMP_DUPLICATE_LIB_OK']='True'

def find_latest_best_model(dir_path):
    """
    查找指定目录中最新的 best 模型文件。
    """
    best_models = []
    # "train"
    for i in range(20):
        if i == 0:
            train_file_path = os.path.join(dir_path, f"train", "weights")
        else:
            train_file_path = os.path.join(dir_path, f"train{i}", "weights")
        
        for root, dirs, files in os.walk(train_file_path):
            for file in files:
                if file.endswith("best.pt"):
                    best_models.append(os.path.join(root, file))
        
    if not best_models:
        # 没有找到任何 best 模型文件，则认为还没开始训练，以YOLO11m作为初始模型开始训练
        return "./yolo11m.pt"
        # 或者没找到，中止训练。
        # raise FileNotFoundError("没有找到 best 模型文件")
    
    # 按文件修改时间排序，取最新的一个
    best_models.sort(key=lambda x: os.path.getmtime(x), reverse=True)
    return best_models[0]

if __name__ == '__main__':
    # 打印当前工作目录
    current_dir = os.getcwd()
    print(f"当前工作目录: {current_dir}")

    # 检查CUDA是否可用
    print(f"CUDA available: {torch.cuda.is_available()}")
    print(f"Number of CUDA devices: {torch.cuda.device_count()}")

    # 创建保存训练结果的目录
    save_dir = os.path.join(current_dir, "runs")
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    # 查找最新的 best 模型文件
    latest_best_model_path = find_latest_best_model(os.path.join(current_dir, "runs", "detect"))
    print(f"找到最新的 best 模型文件: {latest_best_model_path}")

    # Load a model
    # model = YOLO("yolo11n.yaml")  # build a new model from YAML
    # model = YOLO("yolo11m.yaml")  # 换成中杯试试
    # model = YOLO("./yolo11l.pt")  # load a pretrained model (recommended for training)
    # 从本地加载预训练模型
    # model = YOLO("./best.pt")  # 结束之后，可以继续下一轮训练
    # model = YOLO("yolo11n.yaml").load("yolo11n.pt")  # build from YAML and transfer weights
    # Load the latest best model
    model = YOLO(latest_best_model_path)  # load the latest best model

    # Train the model
    results = model.train(data="./sar_aircraft_00.yaml",  # 数据集配置文件路径
                        epochs=40,  # 训练的总轮数
                        imgsz=1024,  # 输入图像的尺寸
                        batch=8,      # 每个批次的大小
                        save_dir=save_dir,  # 将训练结果保存到当前目录
                        # resume=True  # 恢复训练
                        )

    print(f"训练结果已保存到: {save_dir}")